from flask import Flask, request, Response
# from flask_cors import CORS
import dashscope
import os
import json

from flask_cors import CORS
#import imp

app = Flask(__name__)
CORS(app, origins=["http://localhost:8800"])  # # 允许特定源 CORS 支持

# 配置 DashScope API Key
# os.environ["DASHSCOPE_API_KEY"] = "your_api_key_here"
api_key = os.getenv("DASHSCOPE_API_KEY")

@app.route('/api/qwen-plus', methods=['POST'])
def qwen_plus():
    # 获取用户输入
    user_content = request.json.get('content', '')
    # 验证输入
    if not user_content:
        return Response(json.dumps({"error": "Content is required"})), 400
    print(user_content)

    # 构造消息体
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": user_content}
    ]

    try:
        # 调用 DashScope API
        responses = dashscope.Generation.call(
            model="qwen-plus",
            api_key=api_key,
            messages=messages,
            result_format='message',
            stream=True,
            stream_options={"include_usage": True},
            incremental_output=True,
            # response_format={"type": "json_object"}  # 强制 JSON 输出<sup data-index='1'>8</sup>
        )
        

        # 流式返回数据
        def generate():
            for response in responses:
                if response.output.choices and response.output.choices[0].finish_reason != "stop":
                    try:
                        content = response.output.choices[0].message.content
                        yield content
                        #json_content = json.loads(content)
                        #yield json.dumps({"message": json_content}, ensure_ascii=False) + "\n"
                    #except json.JSONDecodeError:
                        #yield json.dumps({"error": "Invalid JSON format from AI"}, ensure_ascii=False) + "\n" 
                    finally:
                        pass    
                else:
                        if response.output.choices is None:
                            content = response.output.usage.total_tokens
                            yield content
                        #yield json.dumps({"error": f"API error: {response.message}"}, ensure_ascii=False) + "\n"
                    

        return Response(generate(), mimetype='application/x-ndjson', headers={'Transfer-Encoding': 'chunked',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive'})
    except Exception as e:
        return Response(json.dumps({"error": str(e)})), 500

if __name__ == '__main__':
    app.run(debug=False,port=8801)#仅用于测试